import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

def complex_potential(Z, a, U, vortices):
    """计算复势函数"""
    F = U * (Z + a**2 / Z)
    for (x_v, y_v, gamma) in vortices:
        z_v = x_v + 1j*y_v
        z_mirror = a**2 / np.conj(z_v)
        term = (Z - z_v) / (Z - z_mirror) * (a / abs(z_v))
        F += 1j * gamma / (2 * np.pi) * np.log(term)
    return F

def visualize_streamlines(a=1.0, U=1.0, vortices=[], res=300):
    """实时流场可视化系统（含动态更新功能）"""
    x = np.linspace(-3*a, 3*a, res)
    y = np.linspace(-3*a, 3*a, res)
    X, Y = np.meshgrid(x, y)
    Z = X + 1j*Y
    
    fig, ax = plt.subplots(figsize=(8, 6))
    circle = plt.Circle((0, 0), a, fc='gray', alpha=0.3)  # 圆柱边界
    ax.add_patch(circle)
    ax.set_xlim(-3*a, 3*a)
    ax.set_ylim(-3*a, 3*a)
    ax.set_aspect('equal')
    plt.title('涡系干扰动态流场')
    plt.xlabel('x')
    plt.ylabel('y')
    
    # 初始流线绘制
    Phi = complex_potential(Z, a, U, vortices)
    psi = np.imag(Phi)
    contour = ax.contour(X, Y, psi, levels=30, cmap='jet')
    plt.colorbar(contour, label='流函数')
    
    # 动态更新函数
    def update(frame):
        ax.clear()
        ax.add_patch(circle)
        ax.set_xlim(-3*a, 3*a)
        ax.set_ylim(-3*a, 3*a)
        ax.set_aspect('equal')
        plt.title(f'涡系干扰动态流场 (t={frame/10:.1f}s)')
        
        # 模拟涡旋运动（示例：简单圆周运动）
        updated_vortices = []
        for (x_v, y_v, gamma) in vortices:
            theta = frame * 0.1  # 旋转角度
            x_rot = x_v * np.cos(theta) - y_v * np.sin(theta)
            y_rot = x_v * np.sin(theta) + y_v * np.cos(theta)
            updated_vortices.append((x_rot, y_rot, gamma))
        
        # 计算更新后的复势
        Phi = complex_potential(Z, a, U, updated_vortices)
        psi = np.imag(Phi)
        # 绘制流线
        ax.contour(X, Y, psi, levels=30, cmap='jet')
        # 标注涡旋位置
        ax.scatter([v[0] for v in updated_vortices], [v[1] for v in updated_vortices], 
                   c='red', marker='o', s=100)
        return ax
    
    # 生成动画（30fps）
    ani = FuncAnimation(fig, update, frames=300, interval=33, blit=False)
    plt.tight_layout()
    return ani

# 示例：含两个旋转涡的动态系统
vortices = [(1.5, 0.5, 2*np.pi), (-1.5, -0.5, -2*np.pi)]
ani = visualize_streamlines(a=1.0, U=1.0, vortices=vortices)
plt.show()